package com.springboot.demo.io;

import java.io.IOException;
import java.net.InetSocketAddress;
import java.nio.ByteBuffer;
import java.nio.channels.SelectionKey;
import java.nio.channels.Selector;
import java.nio.channels.ServerSocketChannel;
import java.nio.channels.SocketChannel;
import java.util.Iterator;
import java.util.Set;

import com.springboot.demo.constant.CommonConstant;

/**
 * 多路复用单线程
 *
 * @author supu
 * @since 2020-11-13 9:32
 **/
public class SocketMultiplexSingleThread {
    int port = 9000;

    /**
     * 多路复用器 linux（select、poll、epoll、kqueue） nginx event{}
     */
    private Selector selector;


    public void initServer() {
        try {
            // 服务端channel，开启socket服务
            ServerSocketChannel socketChannel = ServerSocketChannel.open();
            // 设置非阻塞
            socketChannel.configureBlocking(false);
            // 绑定端口，约等于 listen 状态的 fd4
            socketChannel.bind(new InetSocketAddress(port));

            // 打开多路复用器
            // JDK默认是 epoll，可以通过 -D 设置selector类型
            // epoll 模式下调用了 epoll_create -> fd3
            selector = Selector.open();

            // scoket 服务注册到 selector 上，类型是 accept
            // epoll 模式下调用了 epoll_ctl(fd3, ADD(fd4, EPOLLIN))
            // select、poll 模式下在 JVM 里开辟了一个数组，将 fd4 放到数组中
            socketChannel.register(selector, SelectionKey.OP_ACCEPT);

        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    public void startServer() {
        initServer();

        System.out.println("socketServer start.......");

        // 服务端处理数据
        while (true) {
            try {
                /**
                 * 返回有事件状态的fd数量，此方法会阻塞 调用多路复用器（select、poll or epoll_wait） 1、select、poll ； 调用了 kernel 的
                 * select(fd4)、poll(fd4)，即将所有的fd列表发送给 kernel 2、epoll ； 调用了 kernel 的 epoll_wait() select()
                 * 方法的参数表示阻塞超时时间，可以通过调用 selector.wakeup() 方法 来唤醒阻塞的 selector select() 方法的调用才会触发注册时的 kernel 对
                 * epoll_ctl 的调用
                 */
                int fdNum = selector.select();
                if (fdNum > 0) {
                    // 获取多路复用器上所有有状态的 fd 集合
                    Set<SelectionKey> selectionKeys = selector.selectedKeys();
                    Iterator<SelectionKey> iterator = selectionKeys.iterator();
                    // 循环遍历处理所有有事件的 fd
                    while (iterator.hasNext()) {
                        SelectionKey key = iterator.next();
                        iterator.remove();
                        if (key.isAcceptable()) {
                            // 接收一个新的连接，通过调用 kernel的 accept 方法返回一个新的 fd7，
                            // select、poll 会将fd7保存到 JVM 中（和 listen 的 fd4 放在一起）
                            // epoll 则会调用 kernel的 epoll_ctl 方法注册到 kernel 中
                            acceptHandler(key);
                        } else if (key.isReadable()) {
                            // 处理读事件
                            readHandler(key);
                        }
                    }
                }
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
    }

    private void readHandler(SelectionKey key) {
        // 获取客户端
        SocketChannel client = (SocketChannel) key.channel();
        // 获取绑定的字节 buffer
        ByteBuffer byteBuffer = (ByteBuffer) key.attachment();
        byteBuffer.clear();
        while (true) {
            try {

                // 读取客户端中 buffer 中的数据
                int read = client.read(byteBuffer);
                if (read > 0) {
                    // 有数据
                    byteBuffer.flip();
                    if (byteBuffer.hasRemaining()) {
                        // 读取客户端发来的数据
                        byte[] bytes = new byte[byteBuffer.limit()];
                        byteBuffer.get(bytes);

                        System.out.println(
                                        "client ：" + client.getRemoteAddress() + " send data : " + new String(bytes));
                        // 将数据写回客户端
                        byteBuffer.compact();
                        byteBuffer.put(CommonConstant.getResponseContent().getBytes());
                        byteBuffer.flip();

                        client.write(byteBuffer);
                    }

                    byteBuffer.clear();
                } else if (read == 0) {
                    System.out.println("read over");
                    // 没有数据
                    break;
                } else {
                    System.out.println("client ：" + client.getRemoteAddress() + " closed");
                    client.close();
                    break;
                }
            } catch (IOException e) {
                // 客户端可能已关闭
                try {
                    client.close();
                } catch (IOException ioException) {
                    ioException.printStackTrace();
                }
                break;
            }
        }
    }

    private void acceptHandler(SelectionKey key) {
        try {
            ServerSocketChannel channel = (ServerSocketChannel) key.channel();
            // 调用 kernel的 accept 方法接收客户端并返回一个新的 fd7
            SocketChannel client = channel.accept();
            // 客户端设置非阻塞
            client.configureBlocking(false);

            // 堆外申请 4k 大小的内存 buffer
            ByteBuffer byteBuffer = ByteBuffer.allocateDirect(4096);

            // 客户端 服务注册到 selector 上，类型是 读事件
            // epoll 模式下调用了 epoll_ctl(fd3, ADD(fd4, EPOLLIN))
            // select、poll 模式下在 JVM 里开辟了一个数组，将 fd7 放到数组中
            // 这里为什么不直接注册写事件，是由于 write 事件是通过 send-queue 是否有空间来触发的 ，
            // 一旦注册了 write 事件则会频繁触发（其实好多没有写进来数据，无用功），因此需要先注册 read 事件，
            // read 事件中再注册 write 事件
            client.register(selector, SelectionKey.OP_READ, byteBuffer);

            System.out.println("new client connect : " + client.getRemoteAddress());
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    public static void main(String[] args) {
        SocketMultiplexSingleThread multiplexSingleThread = new SocketMultiplexSingleThread();
        multiplexSingleThread.startServer();
    }
}
